"""
LLM Provider
此模块提供了 LLMClientManager，负责管理所有已配置的 LLM 客户端的生命周期。
"""
from typing import Dict
import httpx
import os

from langchain_core.language_models import BaseChatModel

from axiom_boot.conf.manager import Settings
from axiom_boot.di.decorators import service
from axiom_boot.di.dependency import autowired
from axiom_boot.core.exceptions import FrameworkException
from axiom_boot.logging.setup import get_logger

logger = get_logger(__name__)


@service(name="llm_client_manager")
class LLMClientManager:
    """
    LLM 客户端管理器。

    在应用启动时，根据配置初始化所有可用的 LLM 客户端，
    并提供一个按名称获取特定客户端的接口。
    """
    def __init__(self, settings: Settings = autowired()):
        self._clients: Dict[str, BaseChatModel] = {}
        
        if not settings.langchain.enabled:
            logger.info("LangChain 模块未启用，跳过 LLM 客户端初始化。")
            return

        logger.info("LangChain 模块已启用，正在初始化 LLM 客户端...")
        profiles = settings.langchain.llm_profiles
        if not profiles:
            logger.warning("LangChain 已启用，但未找到任何 'llm_profiles' 配置。")
            return

        for name, profile in profiles.items():
            try:
                client = self._create_client_from_profile(profile)
                self._clients[name] = client
                logger.info(f"LLM 客户端 '{name}' (Provider: {profile.provider}, Model: {profile.model_name}) 已成功初始化。")
            except Exception as e:
                logger.error(f"初始化 LLM 客户端 '{name}' 时发生错误: {e}", exc_info=True)
                # 在启动阶段，这是一个严重问题，可以选择抛出异常中断启动
                raise FrameworkException(f"无法初始化 LLM 客户端 '{name}'。") from e

    def _create_client_from_profile(self, profile) -> BaseChatModel:
        """根据配置档案创建并返回一个 LLM 客户端实例。"""
        provider = profile.provider.lower()

        # --- 通用代理配置 ---
        http_client_args = {}
        
        # 创建可复用的 httpx 客户端
        http_async_client = httpx.AsyncClient(**http_client_args)

        if provider == "openai":
            from langchain_openai import ChatOpenAI
            return ChatOpenAI(
                model=profile.model_name,
                api_key=profile.api_key,
                temperature=profile.temperature,
                http_async_client=http_async_client,
                **profile.extra_args
            )
        elif provider == "google_vertex_ai":
            from langchain_google_vertexai import ChatVertexAI
            from google.oauth2 import service_account

            creds = None
            if profile.service_account_file_path:
                logger.debug(f"从文件加载 Google 凭证: {profile.service_account_file_path}")
                creds = service_account.Credentials.from_service_account_file(
                    profile.service_account_file_path
                )

            return ChatVertexAI(
                model_name=profile.model_name,
                temperature=profile.temperature,
                credentials=creds,
                client_options={"async_client": http_async_client},
                **profile.extra_args
            )
        elif provider == "baidu_qianfan":
            from langchain_community.chat_models import QianfanChatEndpoint
            # 百度千帆需要 AK/SK
            return QianfanChatEndpoint(
                model=profile.model_name,
                qianfan_api_key=profile.api_key,
                qianfan_secret_key=profile.secret_key,
                temperature=profile.temperature,
                http_async_client=http_async_client,
                **profile.extra_args
            )
        elif provider == "dashscope":
            from langchain_community.chat_models import ChatTongyi
            # Dashscope (ChatTongyi) 目前不直接接受 httpx 客户端实例
            # 它内部会自行管理 http 请求
            return ChatTongyi(
                model_name=profile.model_name,
                dashscope_api_key=profile.api_key,
                temperature=profile.temperature,
                **profile.extra_args,
            )
        else:
            # 在抛出异常前确保关闭已创建的 http 客户端，防止资源泄漏
            # 这是一个同步方法，但我们需要调用一个异步关闭。我们使用 asyncio.run 来执行这个操作。
            import asyncio
            try:
                asyncio.run(http_async_client.aclose())
            except RuntimeError:
                 # 如果已经在一个正在运行的事件循环中，run() 会失败。
                 # 在这种情况下，我们假设上层调用者会处理事件循环的关闭。
                 pass
            raise FrameworkException(f"不支持的 LLM provider: '{profile.provider}'")

    def get_client(self, name: str) -> BaseChatModel:
        """
        根据配置档案的名称，获取一个已初始化的 LLM 客户端。

        Args:
            name: 在配置中定义的 LLM profile 名称。

        Returns:
            一个 BaseChatModel 的实例。

        Raises:
            ValueError: 如果找不到具有指定名称的客户端。
        """
        client = self._clients.get(name)
        if not client:
            raise ValueError(f"名为 '{name}' 的 LLM 客户端未找到或未成功初始化。")
        return client 